from align import Cell,Table,SubstMatrix,create_submat
from bio_sm import SubstMatrix_Bio
from Pairwise_Sequence_Alignment import dotplot,indentity

def align_protein():
    seq1,seq2 = 'ANDDR','AARRD'
    sm = SubstMatrix_Bio()
    aligner = Table(seq1,seq2)
    aligner.smith_Waterman(sm,-8)

    for i,pair_seq in enumerate(aligner.recover_align_local()):
        print(f"case {i+1}:")
        for seq in pair_seq:
            print(seq)
    
    print('S matrix:')
    print('\t'+'\t'.join(list(aligner.colum)))
    for line,char in zip(aligner.get_score(),list(aligner.row)):
        print(char+'\t'+ '\t'.join(map(str,line)))
    print('trace-back matrix:')
    print('\t'+'\t'.join(list(aligner.colum)))
    for line,char in zip(aligner.get_trace_back(),list(aligner.row)):
        print(char+'\t'+ '\t'.join(map(str,line)))

def align_dna():
    seq1,seq2 = 'TACT','ACTA'
    sm = SubstMatrix()
    sm.create_submat(3,-1,"ATCG")
    aligner = Table(seq1,seq2)
    aligner.needleman_Wunsch(sm,-3)
    
    for seq in aligner.recover_align():
        print(seq)

    print('S matrix:')
    print('\t'+'\t'.join(list(aligner.colum)))
    for line,char in zip(aligner.get_score(),list(aligner.row)):
        print(char+'\t'+ '\t'.join(map(str,line)))
    print('trace-back matrix:')
    print('\t'+'\t'.join(list(aligner.colum)))
    for line,char in zip(aligner.get_trace_back(),list(aligner.row)):
        print(char+'\t'+ '\t'.join(map(str,line)))

def max_diag(matrix):
    n,m = len(matrix),len(matrix[0])
    start = (0,0)
    size = 0
    for bias in range(-m+1,n):
        x = max(bias,0)
        y = max(-bias,0)
        l = 0
        while x < n and y < m:
            if matrix[x][y] == 1:
                if l == 0:
                    start = (x,y)
                l +=1
            elif l:
                if l > size:
                    size = l
                    l = 0
            # print(matrix[x][y],end='')
            x+=1
            y+=1
        if l > size:
            size = l
    return start,size

def test_max_diag():
    seq1 = "MHQAIFIYQIGYPLKSGYIQSIRSPEYDNW"
    seq2 = "MHQAIFQIYIGYPLPEYDNWKSGYIQSIRS"
    mat = dotplot(seq1,seq2)
    for line in mat:
        print(" ".join(map(str,line)))
    print(max_diag(mat))

class Table_with_ties(Table):

    def get_trace_back(self) -> list[list[int]]:
        T = []
        for i in range(self.shape[0]):
            line = []
            for j in range(self.shape[1]):
                if len(self[i,j].trace)>1:
                    line.append(self[i,j].trace[1:])
                else:
                    line.append([0])
            T.append(line)
        return T
    
    def trace_back(self, start_cell: Cell) -> list[str]:
        seq1,seq2 = '',''
        return self.brach_trace(seq1,seq2,start_cell)
    
    def trace_one_step(self,seq1,seq2,trace,i,j):
        if trace == 1:
            seq1 = self.row[i] + seq1
            seq2 = self.colum[j] + seq2
        elif trace == 2:
            seq1 = '-' + seq1
            seq2 = self.colum[j] + seq2
        elif trace == 3:
            seq1 = self.row[i] + seq1
            seq2 = '-' + seq2
        else:
            raise ValueError(f'error trace {trace}')
        return seq1,seq2

    def brach_trace(self,seq1,seq2,cell:Cell):
        if cell.trace[-1] == 0:
            return  [[seq1,seq2]]
        result = []
        i,j = cell.ij
        if len(cell.precursor) > 1:
            for trace,cell in zip(cell.trace[1:],cell.precursor):
                s1,s2 = self.trace_one_step(seq1,seq2,trace,i,j)
                sub_result = self.brach_trace(s1,s2,cell)
                result.extend(sub_result)
        else:
            trace = cell.trace[-1]
            seq1,seq2 = self.trace_one_step(seq1,seq2,trace,i,j)
            sub_result = self.brach_trace(seq1,seq2,cell.precursor[-1])
            result.extend(sub_result)
        return result
    
    def recover_align_local(self) -> list[list[str]]:
        result = list()
        for start_cell in self.max_cells:
            result.extend(self.trace_back(start_cell))
        return result

def test_needleman_Wunsch_with_ties():
    seq1 = "AAATTAAAGCCC"
    seq2 = "AAATAAAGGCCC"
    gap = -1

    sm = SubstMatrix()
    sm.create_submat(2,-1,"ATCG")
    
    aligner = Table_with_ties(seq1,seq2)
    aligner.needleman_Wunsch(sm,gap)
    for index,seqs in enumerate(aligner.recover_align()):
        print(f"case {index+1}:")
        print(f"{seqs[0]}\n{seqs[1]}")

def search_similar_sequence(suc_seqs,tag_seqs):
    result = []
    for suc in suc_seqs:
        result.append(max(enumerate([indentity(suc,tag) for tag in tag_seqs]),
            key= lambda v:v[1])[0])
    return result

def test_search_similar_sequence():
    suc = ["ATCGCATGCGTACG",
           "CGTACGTATCAGCT",
           "TCGCTATCGTAGCA"]
    tag = ["ACATCGCATGCGTACCAGCT",
           "TCGCTATCGTAGCACTAA",
           "TCAGCGTACGTTATCAGCT"]
    print(search_similar_sequence(suc,tag))

def test_smith_Waterman_with_ties():
    seq1 = "PHSWG"
    seq2 = "HGWAG"
    gap = -10

    sm = SubstMatrix_Bio()
    aligner = Table_with_ties(seq1,seq2)
    aligner.smith_Waterman(sm,gap)

    for index,seqs in enumerate(aligner.recover_align_local()):
        print(f"case {index+1}:")
        print(f"{seqs[0]}\n{seqs[1]}")

if __name__ == "__main__":
    print("### ex1 ###")
    align_protein()
    print("### ex2 ###")
    align_dna()
    print("### ex3 ###")
    test_max_diag()
    print("### ex4 ###")
    test_needleman_Wunsch_with_ties()
    print('### ex5 ###')
    test_smith_Waterman_with_ties()
    print('### ex6 ###')
    test_search_similar_sequence()
    pass